#!/usr/bin/env python
# -*- coding: utf-8 -*-
# @Time : 2024/11/8 11:30
# @Author : XXX
# @Site : 
# @File : 训练部分.py
# @Software: PyCharm
import torch
import torch.nn as nn
import torch.optim as optim


# 定义模型
class Net(nn.Module):
    def __init__(self, input_features, output_classes):
        super(Net, self).__init__()
        self.fc1 = nn.Linear(input_features, 5)
        self.fc2 = nn.Linear(5, 5)  # 第一层隐藏层
        self.fc3 = nn.Linear(5,9)  # 第二层隐藏层
        self.fc4 = nn.Linear(9, 1)  # 第三层隐藏层
        self.fc5 = nn.Linear(1, output_classes)  # 输出层

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.relu(self.fc3(x))
        x = torch.relu(self.fc4(x))
        x = self.fc5(x)  # 输出层通常不加激活函数，除非有特定需求
        return x


# 超参数
input_features = 6  # 根据实际输入数据设定
output_classes = 4  # 根据实际输出设定
learning_rate = 0.001  # 学习率
num_epochs = 50  # 训练轮数

# 创建模型实例
model = Net(input_features, output_classes)

# 损失函数和优化器
criterion = nn.MSELoss()  # 均方误差损失函数
optimizer = optim.Adam(model.parameters(), lr=learning_rate)  # Adam优化器

# 假设你已经有了数据集 X 和标签 Y
# X, Y = ...

# 训练模型
for epoch in range(num_epochs):
    # 前向传播
    outputs = model(X)
    loss = criterion(outputs, Y)

    # 反向传播和优化
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

    if (epoch + 1) % 1 == 0:
        print(f'Epoch [{epoch + 1}/{num_epochs}], Loss: {loss.item():.4f}')